import torch

from config import device
from problems.problem import Problem


class ZDT1(Problem):
    name = 'ZDT1'

    def __init__(self, var_dim, obj_dim, max_fun_eval, kwargs: dict):
        low_limit = torch.zeros((1, var_dim), device=device, dtype=float)
        high_limit = torch.ones((1, var_dim), device=device, dtype=float)
        super().__init__(var_dim, 2, low_limit, high_limit, max_fun_eval, kwargs)

    def eval_value(self, x):
        n = x.shape[1]
        f1 = x[:, 0]
        g = 1 + 9 / (n - 1) * torch.sum(x[:, 1:], dim=1)
        h = 1 - torch.sqrt(f1 + 1e-7 / g)
        f2 = g * h
        return torch.cat((torch.unsqueeze(f1, dim=1), torch.unsqueeze(f2, dim=1)), dim=1)

    def get_optimal_solutions(self, size=10000):
        R = torch.zeros((size, 2), device=device, dtype=torch.double)
        R[:, 0] = torch.linspace(0, 1, size, device=device)
        R[:, 1] = 1 - R[:, 0] ** 0.5
        return R
